﻿package me.zhengjie.common.aspect

import com.google.common.collect.ImmutableList
import me.zhengjie.common.annotation.Limit
import me.zhengjie.common.exception.BadRequestException

import me.zhengjie.common.utils.RequestHolder
import me.zhengjie.common.utils.StringUtils
import jakarta.servlet.http.HttpServletRequest
import org.aspectj.lang.ProceedingJoinPoint
import org.aspectj.lang.annotation.Around
import org.aspectj.lang.annotation.Aspect
import org.aspectj.lang.annotation.Pointcut
import org.aspectj.lang.reflect.MethodSignature
import org.slf4j.LoggerFactory
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.data.redis.core.RedisTemplate
import org.springframework.data.redis.core.script.DefaultRedisScript
import org.springframework.data.redis.core.script.RedisScript
import org.springframework.stereotype.Component

/**
 * @author Kuki
 * @Date 2024/11/10
 */
@Aspect
@Component
class LimitAspect() {
    @Autowired
    lateinit var redisTemplate: RedisTemplate<Any, Any>
    @Pointcut("@annotation(me.zhengjie.common.annotation.Limit)")
    fun pointcut() {
    }

    @Around("pointcut()")
    @Throws(Throwable::class)
    fun around(joinPoint: ProceedingJoinPoint): Any {
        val request: HttpServletRequest = RequestHolder.httpServletRequest
        val signature = joinPoint.signature as MethodSignature
        val signatureMethod = signature.method
        val limit: me.zhengjie.common.annotation.Limit = signatureMethod.getAnnotation(me.zhengjie.common.annotation.Limit::class.java)
        val limitType: LimitType = limit.limitType
        var key: String? = limit.key
        if (org.apache.commons.lang3.StringUtils.isEmpty(key)) {
            key = if (limitType === LimitType.IP) {
                StringUtils.getIp(request)
            } else {
                signatureMethod.name
            }
        }
        val keys: ImmutableList<Any> = ImmutableList.of(
            org.apache.commons.lang3.StringUtils.join(
                limit.prefix,
                "_",
                key,
                "_",
                request.requestURI.toString().replace("/", "_")
            )
        )
        val luaScript = buildLuaScript()
        val redisScript: RedisScript<Number> = DefaultRedisScript(luaScript, Number::class.java)
        val count = redisTemplate.execute(redisScript, keys, limit.count, limit.period)
        return if (null != count && count.toInt() <= limit.count) {
            logger.info("第{}次访问key为 {}，描述为 [{}] 的接口", count, keys, limit.name)
            joinPoint.proceed()
        } else {
            throw BadRequestException("访问次数受限制")
        }
    }

    /**
     * 限流脚本
     */
    private fun buildLuaScript(): String {
        return """
             local c
             c = redis.call('get',KEYS[1])
             if c and tonumber(c) > tonumber(ARGV[1]) then
             return c;
             end
             c = redis.call('incr',KEYS[1])
             if tonumber(c) == 1 then
             redis.call('expire',KEYS[1],ARGV[2])
             end
             return c;
             """.trimIndent()
    }

    companion object {
        private val logger = LoggerFactory.getLogger(LimitAspect::class.java)
    }
}
